args.sharding_dims not guarded against None

At line 153, len(args.sharding_dims) is called unconditionally, but args.sharding_dims can be None when the --sharding-dims flag is omitted (since the argument is not marked required=True and uses nargs="+"). This will raise TypeError: object of type 'NoneType' has no len().

The if len(args.sharding_dims) >= 3: block should be nested inside the existing if args.sharding_dims: guard:

Bias converted unconditionally when use_bias=False

When use_bias=False, self.fc1_bias and self.fc2_bias are initialized as plain torch.Tensor objects (not nn.Parameter, see lines 1940 and 1958):

else:
    self.fc1_bias = torch.Tensor().to(dtype=params_dtype, device=device)

Calling _convert_param_to_dtensor_param on them returns nn.Parameter(DTensor.from_local(...)). When this is then assigned back via self.fc1_bias = new_param, PyTorch's Module.__setattr__ will detect the nn.Parameter type and register the bias as a named module parameter, even though biases are disabled. This would pollute model.named_parameters(), the optimizer parameter list, and checkpoint state.

The fix is to guard these two conversions behind if self.use_bias:, following the same pattern already used for layer_norm_bias at line 2091.

Spurious warning when weight_mesh is None

warnings.warn(...) is emitted unconditionally every time set_device_mesh is called, even when weight_mesh=None. The calling code invokes this method whenever tp_mesh is not None or weight_mesh is not None, so a normal call with only tp_mesh provided will generate a misleading warning like "weight_mesh not necessary for DotProductAttention: None".

The warning should only fire when the caller explicitly passes a non-None weight_mesh. The same spurious warning exists in transformer_engine/pytorch/module/layernorm.py (line 171) and transformer_engine/pytorch/module/rmsnorm.py (line 174).

DCP checkpoint functionality is not exercised in the test

The AppState class (lines 42–93) and DCP checkpoint operations (save, load, get_state_dict, set_state_dict) are imported and fully implemented, but the training loop in _train() (lines 480–490) does not call any checkpoint save/load operations. The function ends at line 497 with dist.destroy_process_group() and no checkpoint round-trip.

Since the PR title is "Add DCP compatibility for FSDP2-TP sharding," the checkpoint functionality is the headline feature. Without an actual save/load call in the test, neither the AppState.state_dict() eviction logic nor the set_state_dict(strict=False) reload path is validated.

Recommendation: Add a checkpoint save/load round-trip after the training loop (before dist.destroy_process_group()) to exercise the DCP functionality, or explicitly note in the test docstring that DCP round-trip testing is deferred to integration tests.

Working on it! + GroupedLinear test case.

args.sharding_dims can be None when --sharding-dims is not passed (the argument uses nargs="+" without required=True). The f-string on this line iterates over it directly, which will raise TypeError: 'NoneType' object is not iterable in that case.

dtensor_member_param.placements was assigned in set_device_mesh relative to the 2-D weight shape (out_features, in_features):

• Column-parallel → Shard(dim=0) (shard over out_features)
• Row-parallel → Shard(dim=1) (shard over in_features)

But the global shape here is the 3-D tensor (num_gemms, out_features, in_features). Reusing the same placements verbatim means:

• Shard(dim=0) now refers to the num_gemms axis — wrong, each TP rank holds all gemms.
• Shard(dim=1) would refer to out_features when the row-parallel split is actually on in_features (dim=2).

As a result, when DCP reconstructs the full checkpoint from the local shards it will use the wrong axis, producing silently corrupted weight tensors.

The shard dimensions need to be incremented by 1 to account for the prepended num_gemms dimension:

from torch.distributed.tensor import Shard as _Shard, Replicate as _Replicate
adjusted_placements = tuple(
    _Shard(p.dim + 1) if isinstance(p, _Shard) else p
    for p in dtensor_member_param.placements
)
grouped_param = _convert_param_to_dtensor_param(
    grouped_param,
    device_mesh=dtensor_member_param.device_mesh,
    placements=adjusted_placements,
    shape=(self.num_gemms, self.out_features, self.in_features),
    stride=None,
)

Okay, this is an impressive catch. I wrote this code too quickly. Will fix!